#!/usr/bin/env python
# Project       PdrUtil
# file name:    pdrutil.py
# Version:      0.1
#
# DESCRIPTION
# Manipulate and create pdr files used by Amazon Kindle devices.
#
# Developer:
# Donald Munro <donaldmunro at gmail dot com>
#
# License:
#         GNU General Public License (GPLv2)
#         For more info see http://www.gnu.org/licenses/gpl-2.0.html

import sys
import os
import optparse
from optparse import OptionParser
import traceback
import tempfile
import shutil
import subprocess
import errno

try:
   from construct import Struct, Container, Array, Embed, UBInt8, UBInt16, UBInt32, NFloat64, PascalString
except:
   sys.stderr.write("""
ERROR: The construct Python package is required (see http://construct.wikispaces.com/) 
or download from http://pypi.python.org/pypi/construct
or install using 
pip install construct 
if you have pip installed
   """ + os.linesep)
   sys.exit(1)

COMMAND = None
bookmark_t = Struct("bookmarks", UBInt8('bookmark_type'), UBInt32('bookmark_page'), PascalString("bookmark_title", length_field = UBInt16("length")))
mark_t = Struct("mark", UBInt8('mark_type'), UBInt32('mark_page'), PascalString("mark_title", length_field = UBInt16("length")),
                PascalString("location", length_field = UBInt16("length")), Array(4, UBInt8("mark_padding")),
                NFloat64("x"), NFloat64("y"))
marking_t = Struct("markings", Embed(Array(2, mark_t)), Array(2, UBInt8("marking_padding"))) 
comment_t = Struct("comments", UBInt8('comment_type'), UBInt32('comment_page'), PascalString("comment_title", length_field = UBInt16("length")),
                   NFloat64("x"), NFloat64("y"), PascalString("location", length_field = UBInt16("length")),
                   PascalString("content", length_field = UBInt16("length")))
pdr_t = Struct("pdr", UBInt32('header'), UBInt8('header_padding'), UBInt32('pageno'),
               UBInt32('nbookmarks'), Array(lambda ctx: ctx["nbookmarks"], bookmark_t),
#               OptionalGreedyRange(UBInt8("ignored")))
               Array(20, UBInt8("padding")),
               UBInt32('nmarkings'), Array(lambda ctx: ctx["nmarkings"], marking_t), 
               UBInt32('ncomments'), Array(lambda ctx: ctx["ncomments"], comment_t),
               Array(4, UBInt8("trailer")))

def main(argv=None):
   global COMMAND, HAS_PDFTK
   if argv is None:
      argv = sys.argv
   usage = """
pdrutil [options] command 
where command can be one of:
read    = Read and decode contents of a pdr file
create  = Create an empty pdr file

add     = Add new bookmarks from a file
replace = Replace bookmarks with new ones from a file
  For the above 2 commands the file to read bookmarks is specified using -I or bookmarkfile
  and the format is multiple lines comprising pageno="Title"eg
  1="Chapter 1"
  2="Section 1.1"
extract = Extract bookmarks from an existing pdr to a text file in the format above.  
pdfextract = Extract bookmarks from a pdf to a text file in the format above.  
  ..     
pdfadd  = Add new bookmarks from the TOC of a pdf file
pdfrep  = Replace bookmarks from the TOC of a pdf file
   For the above 2 commands the file to read bookmarks is specified using -I or bookmarkfile
   and must be a pdf file
sort    = Sort bookmarks on page number

Examples: Read a pdr: pdrutil.py -i Moby-Dick.pdr read
          Create a empty pdr: pdrutil.py -o Lady-Chatterlys-Lover.pdr  create
          Add to a pdr from a text file: pdrutil.py -i Mythical-Man-Month.pdr -o Mythical-Man-Month-out.pdr -B bookmarks.txt add
          Add to a pdr from a pdf file: pdrutil.py -i Mythical-Man-Month.pdr -o Mythical-Man-Month-out.pdr -B Mythical-Man-Month.pdf pdfadd
          Extract from a pdf to a text bookmark file: pdrutil.py -i  Mythical-Man-Month.pdf -o Mythical-Man-Month.txt pdfextract
   """
   parser = OptionParser(usage=usage)
   commands = { 'read' : read, 'create' : create,  'add' : text_add_replace_bookmarks, 'replace' : text_add_replace_bookmarks,
                'sort' : sort_bookmarks, 'extract' : extract_bookmarks, 'pdfextract' : pdf_extract_bookmarks,
                'pdfadd' : pdf_add_replace_bookmarks, 'pdfrep' : pdf_add_replace_bookmarks }
   parser.add_option("-o", '--output',  dest="output", help="Output to pdr file name", default=None)   
   parser.add_option("-O", '--overwrite',  action="store_true", dest="overwrite", default=False,
                     help="Overwrite input pdr file")
   parser.add_option("-i", '--input',  dest="input",  default=None,\
                     help="Input file (if not specified the standard input is assumed)")
   parser.add_option("-t", '--pdftk',  dest="pdftk",  default=None,\
                     help="Path to pdftk if not in PATH")
   parser.add_option("-l", '--level',  dest="level",  default=99999, type="int",\
                     help="When processing PDF bookmarks don't process bookmarks levels higher than this")
   bookmarks_options = optparse.OptionGroup(parser, 'Options for Add/Replace bookmarks')
   bookmarks_options.add_option("-B", '--bookmarkfile',  dest="bookmarkfile", default=None,
                                help="Input file for bookmarks (A text file")
   parser.add_option_group(bookmarks_options)
   (options, args) = parser.parse_args()
   HAS_PDFTK = False   
   if options.pdftk:
      if os.path.isfile(options.pdftk) and os.access(options.pdftk, os.X_OK):
         HAS_PDFTK = True
      else:
         sys.stderr.write('Invalid pdfrk executable %s%s' % (options.pdftk, os.linesep))
         sys.exit(1)
   else:
      if os.name == 'posix':
         program = 'pdftk'
      else:
         program = 'pdftk.exe'
      for path in os.environ["PATH"].split(os.pathsep):
         path = path.strip('"')
         prog = os.path.join(path, program)
         if os.path.isfile(prog) and os.access(prog, os.X_OK):
            HAS_PDFTK = True
            options.pdftk = prog
            break
      if not HAS_PDFTK:
         sys.stderr.write("""
   WARNING: pdftk is required to process PDF files (see http://www.pdflabs.com/tools/pdftk-the-pdf-toolkit/) 
   or download from http://pypi.python.org/pypi/PyPdf2/
   or install using 
   pip install PyPdf2
   if you have pip installed
      """ + os.linesep)
      
   method = None
   if len(args) == 0:
      print('WARNING: COMMAND not specified, assuming read (try -h or --help for invocation format)')
      COMMAND = 'read'
   else:
      COMMAND = args[0].strip().lower()
   method = commands.get(COMMAND, None) 
   if method is None:
      sys.stderr.write('Invalid COMMAND %s%s' % (args[0], os.linesep))
      parser.print_help()
      return 1   
         
   f = None
   try:
      filename = options.input
      if filename is None:
         f = sys.stdin
      else:
         if not os.path.exists(filename):
            sys.stderr.write('File %s not found%s' % (os.path.abspath(filename), os.linesep))
            return 1
         output_filename = options.output
         if not output_filename is None and os.path.exists(output_filename) and os.path.samefile(filename, output_filename):
            sys.stderr.write('Output file %s is the same as input file %s. Specify -O --overwrite instead if you wish to overwrite the input file with the result%s' %\
                             (os.path.abspath(filename), os.path.abspath(output_filename), os.linesep))
            return 1
         f = open(filename, 'r')
      return method(f, options)
   except Exception, ee:
      try:         
         traceback.print_exc(file=sys.stderr)
         sys.stderr.write(str(ee.message) + os.linesep)
      except:
         sys.stderr.write(str(ee.message) + os.linesep)
      return 1
   finally:
      if not f is None and f != sys.stdin:
         try:
            f.close()
         except:
            pass
   
def read(f, options):   
   global pdr_t
   c = pdr_t.parse_stream(f)   
   filename = options.input
   if not filename is None: 
      print('File: ' + os.path.abspath(filename))
   pageno = c.get('pageno', 'N/A')
   print('Current Page: ' + str(pageno))
   bookmarks = c.get('bookmarks', [])
   if len(bookmarks) > 0:
      print('Bookmarks:')
      for bookmark in bookmarks:
         print('   %-40s %s' % (bookmark.get('bookmark_title', 'N/A'),bookmark.get('bookmark_page', 'N/A')))
   markings = c.get('markings', [])
   if len(markings) > 0:
      print('Markings:')
      for marking in markings:
         print(marking)
   comments = c.get('comments', [])
   if len(comments) > 0:
      print('Comments:')
      for comment in comments:
         print(comment)
   return 0

def create(f, options):
   output_filename = options.output
   print 'create ' + output_filename 
   if output_filename is None:
      sys.stderr.write('Please specify an output file using the -o --output option.' + os.linesep)
      return 1      
   if not os.path.isdir(os.path.dirname(os.path.realpath(output_filename))):
      sys.stderr.write('Directory output file %s does not exist or is not a directory.%s' % (os.path.realpath(output_filename), os.linesep))
      return 1
   if os.path.exists(output_filename):
      sys.stderr.write('Output file %s exists.%s' % (os.path.realpath(output_filename), os.linesep))
      return 1
   c = Container(header=3735931579, header_padding=4, pageno=1, nbookmarks=0, ncomments=0, nmarkings=0,
                 bookmarks=[], comments=[], markings=[],\
                 padding=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 223, 0, 0, 0, 0, 0, 0, 2, 48],
                 trailer=[0, 0, 0, 2])
   isok = False
   try:
      outputfile = open(output_filename, 'w')
      pdr_t.build_stream(c, outputfile)
      isok = True
      print('Created '+ output_filename)
   finally:
      if not outputfile is None:
         outputfile.close()
      if not isok and os.path.exists(output_filename):
         try:
            os.remove(output_filename)
         except:
            pass
                        
def text_add_replace_bookmarks(f, options):
   global pdr_t, bookmark_t
   global COMMAND
   if not options.output is None and options.overwrite:
      sys.stderr.write('Cannot specify both output file (-o --output) and overwrite (-O --overwrite)' + os.linesep)
      return 1
   bookmark_filename = options.bookmarkfile
   if bookmark_filename is None:
      sys.stderr.write('Please specify a text file containing the bookmarks using the -B or --bookmarkfile option' + os.linesep)
      return 1
   if not os.path.exists(bookmark_filename):
      sys.stderr.write('Bookmark file %s not found%s' % (os.path.abspath(bookmark_filename), os.linesep))
      return 1
   bookmark_file = None   
   output_filename = options.output
   if output_filename is None and not options.overwrite:
      sys.stderr.write('Please specify an output file using the -o --output option or specify overwrite mode to overwrite the input file' + os.linesep)
      return 1      
   if not output_filename is None and not os.path.isdir(os.path.dirname(os.path.realpath(output_filename))):
      sys.stderr.write('Directory output file %s does not exist or is not a directory.%s' % (os.path.realpath(output_filename), os.linesep))
      return 1
   outputfile = None
   is_add = (COMMAND == 'add')
   isok = False
   try:
      bookmark_file = open(bookmark_filename, 'r')
      if options.overwrite:
         (fd, output_filename) = tempfile.mkstemp(".pdr", "txt")
         outputfile = os.fdopen(fd, 'w')
      else:
         outputfile = open(output_filename, 'w')
      c = pdr_t.parse_stream(f)
      bookmarks = c.get('bookmarks', [])      
      if not is_add:
         bookmarks = []
      line = bookmark_file.readline()
      while line != '':
         line = line.strip()
         if line == '' or line.startswith('#'):
            line = bookmark_file.readline()
            continue
         p = line.find('=')
         if p < 0:
            line = bookmark_file.readline()
            continue
         page = line[0:p].strip()
         if page == '':
            line = bookmark_file.readline()
            continue
         title = None
         try:
            title = line[p+1:].strip()
            if title[0] == '"':
               title = title[1:]
            if title[len(title)-1] == '"':
               title = title[0:len(title)-1]
         except:
            line = bookmark_file.readline()
            continue
         if title is None:
            title = page
#         bookmarks.append(bookmark_t.build(Container(bookmark_type = 0, bookmark_page = int(page), bookmark_title = title)))
         bookmarks.append(Container(bookmark_type = 0, bookmark_page = int(page), bookmark_title = title))
         line = bookmark_file.readline()
      c['bookmarks'] = bookmarks
      c['nbookmarks'] = len(bookmarks)
      pdr_t.build_stream(c, outputfile)
      isok = True
   finally:
      if not outputfile is None:
         outputfile.close()
      if not isok and os.path.exists(output_filename):
         try:
            os.remove(output_filename)
         except:
            pass
      if not bookmark_file is None:
         bookmark_file.close()         
      if options.overwrite:
         if isok and os.path.exists(output_filename):
            filename = options.input
            if not filename is None:
               if os.path.exists(filename):
                  os.remove(filename)
               shutil.move(output_filename, filename)
               
def sort_bookmarks(f, options):
   global pdr_t, bookmark_t
   global COMMAND
   if not options.output is None and options.overwrite:
      sys.stderr.write('Cannot specify both output file (-o --output) and overwrite (-O --overwrite)' + os.linesep)
      return 1
   output_filename = options.output
   if output_filename is None and not options.overwrite:
      sys.stderr.write('Please specify an output file using the -o --output option or specify overwrite mode to overwrite the input file' + os.linesep)
      return 1      
   if not output_filename is None and not os.path.isdir(os.path.dirname(os.path.realpath(output_filename))):
      sys.stderr.write('Directory output file %s does not exist or is not a directory.%s' % (os.path.realpath(output_filename), os.linesep))
      return 1
   outputfile = None
   isok = False
   try:
      if options.overwrite:
         (outputfile, output_filename) = tempfile.mkstemp(".tmp", "REQ")
      else:
         outputfile = open(output_filename, 'w')
      c = pdr_t.parse_stream(f)
      bookmarks = c.get('bookmarks', [])
      bookmarks.sort(key=lambda bookmark: bookmark.get('bookmark_page', 0))
      c['bookmarks'] = bookmarks
      pdr_t.build_stream(c, outputfile)
      isok = True
   finally:
      if not outputfile is None:
         outputfile.close()
      if not isok and os.path.exists(output_filename):
         try:
            os.remove(output_filename)
         except:
            pass
      if options.overwrite:
         if isok and os.path.exists(output_filename):
            filename = options.input
            if not filename is None:
               if os.path.exists(filename):
                  os.remove(filename)
               shutil.move(output_filename, filename)
               
def extract_bookmarks(f, options):
   global pdr_t, bookmark_t
   global COMMAND
   bookmark_filename = options.bookmarkfile
   if bookmark_filename is None:
      sys.stderr.write('Please specify a text file to which to extract the bookmarks using the -B or --bookmarkfile option' + os.linesep)
      return 1
   bookmark_file = None   
   try:
      bookmark_file = open(bookmark_filename, 'w')
      c = pdr_t.parse_stream(f)
      bookmarks = c.get('bookmarks', [])
      bookmarks.sort(key=lambda bookmark: bookmark.get('bookmark_page', 0))
      for bookmark in bookmarks:
         page = bookmark.get('bookmark_page', 0)
         title = bookmark.get('bookmark_title', "").strip()
         if page == 0 and title == '':
            continue
         bookmark_file.write('%d="%s"%s' % (page, title, os.linesep))
   finally:
      if not bookmark_file is None:
         bookmark_file.close()         

def pdf_extract_bookmarks(f, options):
   global pdr_t, bookmark_t
   global COMMAND, HAS_PDFTK
   if not HAS_PDFTK:
      sys.stderr.write('Please install pdftk to process PDF files for bookmarks.' + os.linesep)
      return 1
   bookmark_filename = ''
   bookmark_filename = options.input
   if bookmark_filename is None or not bookmark_filename.lower().endswith(".pdf"):
      sys.stderr.write('Please specify a pdf file containing a usable outline using the -i or --input option' + os.linesep)
      return 1
   if not os.path.exists(bookmark_filename):
      sys.stderr.write('Pdf file %s not found%s' % (os.path.abspath(bookmark_filename), os.linesep))
      return 1
   output_file = None      
   try:
      ll = extract_pdf_bookmarks(bookmark_filename, options)
      if ll is None or len(ll) == 0:
         sys.stderr.write('PDF file did not contain a usable outline.%s' % (os.linesep,))
         return 1
      output_filename = options.output
      output_file = open(output_filename, "w")
      for t in ll:
         page = t[0]
         title = None
         if len(t) > 1:
            title = t[1]
         if title is None:
            title = page
         if page == 0 and title == '':
            continue
         try:
            output_file.write('%d="%s"%s' % (page, title, os.linesep))
         except:
            pass
   finally:
      if not output_file is None:
         output_file.close()         

def pdf_add_replace_bookmarks(f, options):
   global pdr_t, bookmark_t, HAS_PDFTK
   global COMMAND
   if not HAS_PDFTK:
      sys.stderr.write('Please install pdftk to process PDF files for bookmarks.' + os.linesep)
      return 1
   if not options.output is None and options.overwrite:
      sys.stderr.write('Cannot specify both output file (-o --output) and overwrite (-O --overwrite)' + os.linesep)
      return 1
   bookmark_filename = options.bookmarkfile
   if bookmark_filename is None:
      sys.stderr.write('Please specify a pdf file containing a usable outline using the -B or --bookmarkfile option' + os.linesep)
      return 1
   if not os.path.exists(bookmark_filename):
      sys.stderr.write('Bookmark file %s not found%s' % (os.path.abspath(bookmark_filename), os.linesep))
      return 1
   bookmark_file = None   
   output_filename = options.output
   if output_filename is None and not options.overwrite:
      sys.stderr.write('Please specify an output file using the -o --output option or specify overwrite mode to overwrite the input file' + os.linesep)
      return 1      
   if not output_filename is None and not os.path.isdir(os.path.dirname(os.path.realpath(output_filename))):
      sys.stderr.write('Directory output file %s does not exist or is not a directory.%s' % (os.path.realpath(output_filename), os.linesep))
      return 1
   outputfile = None
   is_add = (COMMAND == 'pdfadd')
   isok = False
   try:
      ll = extract_pdf_bookmarks(bookmark_filename, options)
      if ll is None or len(ll) == 0:
         sys.stderr.write('PDF file %s did not contain a usable outline.%s' % (bookmark_filename,os.linesep))
         return 1
      if options.overwrite:
         (fd, output_filename) = tempfile.mkstemp(".pdr", "pdf")
         outputfile = os.fdopen(fd, 'w')
      else:
         outputfile = open(output_filename, 'w')
      c = pdr_t.parse_stream(f)
      bookmarks = c.get('bookmarks', [])      
      if not is_add:
         bookmarks = []
      for t in ll:
         page = str(t[0])
         if page == '':
            continue
         title = None
         if len(t) > 1:
            title = t[1]
         if title is None:
            title = page
         bookmarks.append(Container(bookmark_type = 0, bookmark_page = int(page), bookmark_title = title))
      c['bookmarks'] = bookmarks
      c['nbookmarks'] = len(bookmarks)
      pdr_t.build_stream(c, outputfile)
      isok = True
   finally:
      if not outputfile is None:
         outputfile.close()
      if not isok and os.path.exists(output_filename):
         try:
            os.remove(output_filename)
         except:
            pass
      if not bookmark_file is None:
         bookmark_file.close()         
      if options.overwrite:
         if isok and os.path.exists(output_filename):
            filename = options.input
            if not filename is None:
               if os.path.exists(filename):
                  os.remove(filename)
               shutil.move(output_filename, filename)

def extract_pdf_bookmarks(pdffilename, options):
   ll = []
   def do_append(pg, title, level, ll):
      page = -1;
      if not pg is None:
         try:
            page = int(pg.strip())
         except:
            page = -1;
      if title is None:
         title = pg
      else:
         title = title.strip()
      if level <= options.level and page >= 0:
#         print((page, title, type(page)))                               
         ll.append((page, title))   
   pdftk = subprocess.Popen([options.pdftk, pdffilename, 'dump_data'],  stdout=subprocess.PIPE,)
   title = None
   pg = None
   level = 2   
   line = pdftk.stdout.readline()
   while not line is None and line != '':      
      line = line.strip() 
      if line.startswith('BookmarkBegin'):
         if not title is None or not pg is None:
            do_append(pg, title, level, ll)
            title = None
            pg = None
      elif line.startswith('BookmarkTitle:') and len(line) > 14:
         title = line[14:]
      elif line.startswith('BookmarkLevel:')  and len(line) > 14:
         try:
            level = int(line[14:].strip())
         except:
            level = 2
      elif line.startswith('BookmarkPageNumber:') and len(line) > 19:
         pg = line[19:]
      line = pdftk.stdout.readline()
   if not title is None or not pg is None:
      do_append(pg, title, level, ll)
   pdftk.wait()   
   pdftk.stdout.close()
   return ll
         
if __name__ == "__main__":
   sys.exit(main())
